#!/usr/bin/python3
import os
import numpy as np
import pandas as pd

face_attr_list_file = r'/rootfs/media/kasim/Data1/data/VideoFaceImage/extract_image_list_72.attr.txt'
face_filter_list_file = r'/rootfs/media/kasim/Data1/data/VideoFaceImage/filter3/FaceRecSel.txt'


def main():
    filter_file_names = set()
    file_count = 0
    faceids = set()
    with open(face_filter_list_file, 'r') as file:
        for line in file.readlines():
            lines = line.split()
            file_name = lines[-2]
            filter_file_names.add(file_name)
            file_count += 1
            if file_count % 10000 == 0:
                print('Process Read File Count: {}'.format(file_count))
            faceids.add(lines[0]+lines[1])
        print('Process Read File Count: {}'.format(file_count))
        print('Face ID Count:', len(faceids))

    # file_name;bin_pitch;bin_roll;bin_yaw;pitch;roll;yaw;quality;mask
    attr_cvs = pd.read_csv(face_attr_list_file, sep=';')
    file_names = attr_cvs['file_name']
    pitchs = attr_cvs['pitch']
    rolls = attr_cvs['roll']
    yaws = attr_cvs['yaw']
    qualitys = attr_cvs['quality']
    file_count = 0
    out_file_path = os.path.join(os.path.dirname(face_filter_list_file), 'FaceAttrs.txt')
    out_file = open(out_file_path, 'w')
    for file_name, quality, pitch, roll, yaw in zip(file_names, qualitys, pitchs, rolls, yaws):
        if file_name in filter_file_names:
            quality = list(map(float, quality.split(',')))
            # quality = np.array(quality, dtype=np.float32)
            # max_idx = np.argmax(quality)
            # score = quality[max_idx]
            # attrs_info = '{};{:.03f},{:.03f},{:.03f};{},{:.03f}\n'.format(file_name, pitch, roll, yaw, max_idx, score)
            score = 0.3333333333*quality[2] + 0.66666666*quality[1] + quality[0]
            attrs_info = '{};{:.03f},{:.03f},{:.03f};{:.03f}\n'.format(file_name, pitch, roll, yaw, score)
            out_file.write(attrs_info)
            file_count += 1
            if file_count % 10000 == 0:
                print('Process Filter File Count: {}'.format(file_count))
    out_file.close()
    print('Process Filter File Count: {}'.format(file_count))
    os.system('chmod a+wr {}'.format(out_file_path))


if __name__ == '__main__':
    main()
